/*
 * Copyright (c) 2021-2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "quickener.h"
#include "libpandafile/debug_data_accessor-inl.h"
#include "libpandafile/line_number_program.h"
#include "libpandabase/utils/utf.h"
#include <cstdint>
#include <iostream>

namespace ark::quick {

panda_file::DebugInfoItem *Quickener::CreateDebugInfoItem(panda_file::File::EntityId debug_info_id)
{
    auto it = ids_done.find(debug_info_id);
    if (it != ids_done.end()) {
        return static_cast<panda_file::DebugInfoItem *>(it->second);
    }

    auto *lnp_item = container_->CreateLineNumberProgramItem();
    auto *debug_info_item = container_->CreateItem<panda_file::DebugInfoItem>(lnp_item);
    ids_done.insert({debug_info_id, static_cast<panda_file::BaseItem *>(debug_info_item)});

    panda_file::DebugInfoDataAccessor debug_acc(*file_, debug_info_id);

    debug_info_item->SetLineNumber(debug_acc.GetLineStart());
    debug_acc.EnumerateParameters([&](panda_file::File::EntityId param_id) {
        auto data = file_->GetStringData(param_id);
        std::string item_str(utf::Mutf8AsCString(data.data));
        auto *string_item = container_->GetOrCreateStringItem(item_str);
        debug_info_item->AddParameter(string_item);
    });

    return debug_info_item;
}

void Quickener::UpdateDebugInfo(panda_file::DebugInfoItem *debug_info_item, panda_file::File::EntityId debug_info_id)
{
    auto *lnp_item = debug_info_item->GetLineNumberProgram();
    panda_file::DebugInfoDataAccessor debug_acc(*file_, debug_info_id);
    const uint8_t *program = debug_acc.GetLineNumberProgram();
    auto size = file_->GetSpanFromId(file_->GetIdFromPointer(program)).size();
    auto opcode_sp = Span(program, size);

    size_t i = 0;
    panda_file::LineNumberProgramItem::Opcode opcode;
    panda_file::LineProgramState state(*file_, panda_file::File::EntityId(0), debug_acc.GetLineStart(),
                                       debug_acc.GetConstantPool());
    while ((opcode = panda_file::LineNumberProgramItem::Opcode(opcode_sp[i++])) != panda_file::LineNumberProgramItem::Opcode::END_SEQUENCE) {
        switch (opcode) {
            case panda_file::LineNumberProgramItem::Opcode::ADVANCE_PC: {
                lnp_item->EmitAdvancePc(debug_info_item->GetConstantPool(), state.ReadULeb128());
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::ADVANCE_LINE: {
                lnp_item->EmitAdvanceLine(debug_info_item->GetConstantPool(), state.ReadSLeb128());
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::START_LOCAL: {
                auto [reg_number, n, is_full] = leb128::DecodeSigned<int32_t>(&opcode_sp[i]);
                LOG_IF(!is_full, FATAL, COMMON) << "Cannot read a register number";
                i += n;

                auto name_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string name = utf::Mutf8AsCString(file_->GetStringData(name_id).data);
                auto *name_item = container_->GetOrCreateStringItem(name);

                auto type_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string type_name = utf::Mutf8AsCString(file_->GetStringData(type_id).data);
                auto *type_item = file_->IsExternal(type_id)
                                      ? static_cast<panda_file::BaseClassItem *>(container_->GetOrCreateForeignClassItem(type_name))
                                      : static_cast<panda_file::BaseClassItem *>(container_->GetOrCreateClassItem(type_name));

                lnp_item->EmitStartLocal(debug_info_item->GetConstantPool(), reg_number, name_item,
                                         type_item->GetNameItem());
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::START_LOCAL_EXTENDED: {
                auto [reg_number, n, is_full] = leb128::DecodeSigned<int32_t>(&opcode_sp[i]);
                LOG_IF(!is_full, FATAL, COMMON) << "Cannot read a register number";
                i += n;

                auto name_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string name = utf::Mutf8AsCString(file_->GetStringData(name_id).data);
                auto *name_item = container_->GetOrCreateStringItem(name);

                auto type_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string type_name = utf::Mutf8AsCString(file_->GetStringData(type_id).data);
                auto *type_item = file_->IsExternal(type_id)
                                      ? static_cast<panda_file::BaseClassItem *>(container_->GetOrCreateForeignClassItem(type_name))
                                      : static_cast<panda_file::BaseClassItem *>(container_->GetOrCreateClassItem(type_name));

                auto type_signature_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string type_signature = utf::Mutf8AsCString(file_->GetStringData(type_signature_id).data);
                auto *type_signature_item = container_->GetOrCreateStringItem(type_signature);

                lnp_item->EmitStartLocalExtended(debug_info_item->GetConstantPool(), reg_number, name_item,
                                                 type_item->GetNameItem(), type_signature_item);
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::END_LOCAL: {
                auto [reg_number, n, is_full] = leb128::DecodeSigned<int32_t>(&opcode_sp[i]);
                LOG_IF(!is_full, FATAL, COMMON) << "Cannot read a register number";
                i += n;

                lnp_item->EmitEndLocal(reg_number);
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::RESTART_LOCAL: {
                auto [reg_number, n, is_full] = leb128::DecodeSigned<int32_t>(&opcode_sp[i]);
                LOG_IF(!is_full, FATAL, COMMON) << "Cannot read a register number";
                i += n;

                lnp_item->EmitRestartLocal(reg_number);
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::SET_PROLOGUE_END: {
                lnp_item->EmitPrologueEnd();
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::SET_EPILOGUE_BEGIN: {
                lnp_item->EmitEpilogueBegin();
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::SET_FILE: {
                auto source_file_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string source_file = utf::Mutf8AsCString(file_->GetStringData(source_file_id).data);
                auto *source_file_item = container_->GetOrCreateStringItem(source_file);
                lnp_item->EmitSetFile(debug_info_item->GetConstantPool(), source_file_item);
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::SET_SOURCE_CODE: {
                auto source_code_id = panda_file::File::EntityId(state.ReadULeb128());
                std::string source_code = utf::Mutf8AsCString(file_->GetStringData(source_code_id).data);
                auto *source_code_item = container_->GetOrCreateStringItem(source_code);
                lnp_item->EmitSetFile(debug_info_item->GetConstantPool(), source_code_item);
                break;
            }
            case panda_file::LineNumberProgramItem::Opcode::SET_COLUMN: {
                lnp_item->EmitColumn(debug_info_item->GetConstantPool(), 0, state.ReadULeb128());
                break;
            }
            default: {
                auto opcode_value = static_cast<uint8_t>(opcode);
                auto adjust_opcode = opcode_value - panda_file::LineNumberProgramItem::OPCODE_BASE;
                uint32_t pc_diff = adjust_opcode / panda_file::LineNumberProgramItem::LINE_RANGE;
                int32_t line_diff =
                    adjust_opcode % panda_file::LineNumberProgramItem::LINE_RANGE + panda_file::LineNumberProgramItem::LINE_BASE;
                lnp_item->EmitSpecialOpcode(pc_diff, line_diff);
                break;
            }
        }
    }
    lnp_item->EmitEnd();
}

void Quickener::QuickContainer()
{
    auto *class_map = container_->GetClassMap();

    std::unordered_map<panda_file::DebugInfoItem *, panda_file::File::EntityId> new_info_map;

    for (const auto &it : *class_map) {
        auto *base_class_item = it.second;
        if (base_class_item->IsForeign()) {
            continue;
        }

        auto *class_item = static_cast<panda_file::ClassItem *>(base_class_item);

        const std::unordered_map<uint32_t, uint32_t> *translation_map = nullptr;
        switch (class_item->GetSourceLang()) {
% Panda.quickened_plugins.each_key do |namespace|
            case panda_file::SourceLang::<%= namespace.upcase %>:
                translation_map = &<%= namespace.upcase %>_TRANSLATION_MAP;
                break;
% end
            default:
                continue;
        }

        if(translation_map == nullptr) {
            continue;
        }

        class_item->VisitMethods([&](panda_file::BaseItem *param_item) {
            auto *method_item = static_cast<panda_file::MethodItem *>(param_item);
            auto *code_item = method_item->GetCode();
            if (code_item == nullptr) {
                return true;
            }

            method_item->SetCode(GetQuickenedCode(code_item, translation_map));
            auto* debug_info_item = static_cast<panda_file::BaseItem *>(method_item->GetDebugInfo());
            for(const auto& item : *items_){
                if(item.second == debug_info_item){
                    auto* new_debug_info = CreateDebugInfoItem(item.first);
                    new_info_map.emplace(new_debug_info, item.first);
                    method_item->SetDebugInfo(new_debug_info);
                    debug_info_item->SetNeedsEmit(false);
                    break;
                }
            }

            code_item->SetNeedsEmit(false);

            return true;
        });
    }

    container_->ComputeLayout();

    for(auto& inf : new_info_map){
        if(inf.first != nullptr){
            UpdateDebugInfo(inf.first, inf.second);
        }
    }

    container_->SetQuickened();
    container_->DeduplicateItems(true);
}

panda_file::CodeItem *Quickener::GetQuickenedCode(panda_file::CodeItem *code, const std::unordered_map<uint32_t, uint32_t> *translation_map)
{
    // branch original offset -> branch imm size in bytes
    struct BranchInfo {
        size_t offset;
        size_t imm_size;
    };
    // TODO(nsizov): refactor this function when BytecodeInstr api would be done for quickened instrs
    std::unordered_map<size_t, size_t> offset_map;
    std::vector<BranchInfo> branches;
    std::vector<uint8_t> new_instructions;
    std::vector<ark::panda_file::CodeItem::TryBlock> new_try_blocks;

    size_t offset = 0;
    offset_map.emplace(offset, 0);
    BytecodeInstruction inst(code->GetInstructions()->data());
    while (offset < code->GetCodeSize()) {
        const size_t pref_size = inst.IsPrefixed() ? 2 : 1;

        // save offset to a jump instruction and its imm bytes count
        if (inst.HasFlag(BytecodeInstruction::Flags::JUMP)) {
            branches.emplace_back(BranchInfo {offset, inst.GetSize() - pref_size});
        }

        size_t copy_idx = 0;
        // TODO(nsizov): deduct actual map depending on method sourcelang
        auto it = translation_map->find(static_cast<size_t>(inst.GetOpcode()));
        if (it != translation_map->end()) {
            new_instructions.push_back(it->second);
            copy_idx = pref_size;
        }
        while (copy_idx < inst.GetSize()) {
            new_instructions.push_back(inst.ReadByte(copy_idx++));
        }

        offset += inst.GetSize();
        inst = inst.GetNext();

        offset_map.emplace(offset, new_instructions.size());
    }

    for (auto &block : code->GetTryBlocks()) {
        std::vector<ark::panda_file::CodeItem::CatchBlock> new_catch_blocks;
        for (auto &catch_block : block.GetCatchBlocks())
            {
                auto size = offset_map[catch_block.GetHandlerPc() + catch_block.GetCodeSize()] -
                            offset_map[catch_block.GetHandlerPc()];

                new_catch_blocks.emplace_back(ark::panda_file::CodeItem::CatchBlock(
                    catch_block.GetMethod(), catch_block.GetType(), offset_map[catch_block.GetHandlerPc()], size));
                catch_block.SetNeedsEmit(false);
            }

        auto length = offset_map[block.GetStartPc() + block.GetLength()] - offset_map[block.GetStartPc()];

        block.SetNeedsEmit(false);

        new_try_blocks.emplace_back(
            ark::panda_file::CodeItem::TryBlock(offset_map[block.GetStartPc()], length, new_catch_blocks));
    }

    for (auto &branch : branches) {
        auto it_branch_offset = offset_map.find(branch.offset);
        LOG_IF(it_branch_offset == offset_map.end(), FATAL, QUICKENER) << "Invalid branch offset";

        ssize_t old_jump_offset = 0;
        for (size_t i = branch.imm_size; i > 0; i--) {
            // NOLINTNEXTLINE(hicpp-signed-bitwise)
            old_jump_offset <<= CHAR_BIT;
            // NOLINTNEXTLINE(hicpp-signed-bitwise)
            old_jump_offset |= (new_instructions[it_branch_offset->second + i]);
        }

        // sign extend the offset
        const uint8_t bytes_to_shift = CHAR_BIT * branch.imm_size - 1;
        // NOLINTNEXTLINE(hicpp-signed-bitwise, clang-analyzer-core.UndefinedBinaryOperatorResult)
        if (old_jump_offset >> bytes_to_shift == 1) {
            // NOLINTNEXTLINE(hicpp-signed-bitwise)
            old_jump_offset |= ((~old_jump_offset >> bytes_to_shift) << bytes_to_shift);
        }
        auto branch_target = offset_map.find(branch.offset + old_jump_offset);
        LOG_IF(branch_target == offset_map.end(), FATAL, QUICKENER) << "Invalid jump target offset";

        int64_t new_jump_offset = branch_target->second - it_branch_offset->second;

        static constexpr const uint8_t FIRST_BYTE_MASK = 0xff;
        for (size_t i = 1; i <= branch.imm_size; i++) {
            // NOLINTNEXTLINE(hicpp-signed-bitwise)
            new_instructions[it_branch_offset->second + i] = (new_jump_offset & FIRST_BYTE_MASK);
            // NOLINTNEXTLINE(hicpp-signed-bitwise)
            new_jump_offset >>= CHAR_BIT;
        }
    }

    auto* new_code = container_->CreateItem<panda_file::CodeItem>(code->GetNumVregs(), code->GetNumArgs(),
                                                        std::move(new_instructions));

    for(auto& block : new_try_blocks) {
        new_code->AddTryBlock(block);
    }

    return new_code;
}

}  // namespace ark::quick
